import numpy as np
import matplotlib.pyplot as plt
import math
#Code for figure
#create function
np.random.seed(12)
ss= 10
A = np.random.normal(size=(ss, ss))
b = np.random.normal(size=(ss, 1))
x1 = np.random.normal(size=(ss, 1))
A1 = np.transpose(A)
V1 = np.dot(A1, A)

def f(x1):
    u = np.dot(A, x1) - b
    return np.linalg.norm(u) ** 2
#原函数，输入n维list，输出1维list

def grad(x1):
    D = np.dot(V1, 2 * x1) - np.dot(A1, 2 * b)
    D = D.tolist()
    D = np.array(D)
    return D
#梯度，输入n维list,输出n维list

def compare(x1,x2):
    return x1 if f(x1)<=f(x2) else x2
#比较函数，输入两个n维list,输出一个n维list
def gradw(x,rate,y):
    return [grad(x)[i]+((x[i]-y[i])/math.sqrt(rate)) for i in range(len(grad(x)))]
def Pr(x,y,iter,rate):
    for i in range(iter - 1):
        x = x - np.dot(rate , gradw(x,rate,y))
    return x
def PROX(origin,epochs,rate,p):
    x = np.array(origin)
    y = np.array(origin)
    z = np.array(origin)
    dot_set = [x]
    for i in range(epochs - 1):
        y = (1 - co(p, i)) * x + co(p, i) * z
        x = Pr(x,y,iter=50,rate=rate)
        z = z - (rate / (co(p, i))) * grad(x)
        dot_set.append(x)
    return dot_set
#GNSA，输出epochs*n 维列表
def GD(origin,epochs,rate):
    x= np.array(origin)
    dot_set = [x]
    for i in range(epochs-1):
        x=x-rate*grad(x)
        dot_set.append(x)
    return dot_set
#梯度下降，输入初值，步数，学习率，输出epochs*n 维列表

def co(p,t):
    return p/(t+p)
def NAG(origin,epochs,rate,p):
    x= np.array(origin)
    y= np.array(origin)
    dot_set = [x]
    for i in range(epochs-1):
        z=x
        x = y-np.dot(rate,grad(y))
        y = x+np.dot((i)/(i+4),x-z)
        dot_set.append(x)
    return dot_set
#Nesterov梯度下降，输出epochs*n 维列表
def gradw(x,rate,y):
    return [grad(x)[i]+((x[i]-y[i])/math.sqrt(rate)) for i in range(len(grad(x)))]
def Pr(x,y,iter,rate):
    for i in range(iter - 1):
        x = x - np.dot(rate , gradw(x,rate,y))
    return x
def GNSA(origin,epochs,rate,p):
    x = np.array(origin)
    y = np.array(origin)
    z = np.array(origin)
    dot_set = [x]
    for i in range(epochs - 1):
        y = (1 - co(p, i)) * x + co(p, i) * z
        x = compare(y - rate * grad(y),x- rate*grad(x))
        z = z - (rate / (co(p, i))) * grad(y)
        dot_set.append(x)
    return dot_set
def PROX(origin,epochs,rate,p):
    x = np.array(origin)
    y = np.array(origin)
    z = np.array(origin)
    dot_set = [x]
    for i in range(epochs - 1):
        y = (1 - co(p, i)) * x + co(p, i) * z
        x = Pr(x,y,iter=20,rate=rate)
        z = z - (rate / (co(p, i))) * grad(y)
        dot_set.append(x)
    return dot_set
def pr_track(y):
    return
def pr_fun_value(y,min):
    plt.semilogy([i for i in range(len(y))],[f(y[i]) for i in range(len(y))])
    return
#绘制轨迹图与函数值下降图
#调试
end=1000
or1=np.array(np.ones(ss).reshape([ss,1]))
e1=0.005
#print(grad_descent(origin=[-2,3,5],epochs=end,rate=0.01)[end-1])#梯度下降调试
#print(NSA(origin=[-2,3,5],epochs=end,rate=0.01,p=3)[end-1])#Nesterov梯度下降测试
#print(GNSA(origin=[-2,3,5],epochs=end,rate=0.01,p=3)[end-1])
pr_fun_value(GD(origin=or1,epochs=end,rate=e1),0)
pr_fun_value(NAG(origin=or1,epochs=end,rate=e1,p=3),0)
pr_fun_value(PROX(origin=or1,epochs=end,rate=e1,p=3),0)
plt.legend(loc='best',frameon=True, labels=['GD','NAG','PROX'])
plt.xlabel('iteration')
plt.ylabel('f - f*')
plt.show()